import socket
import struct
import numpy as np
from scipy.io import wavfile
import datetime
import os
import time
import wave

# ================= 配置 =================
HOST = '0.0.0.0'  # 监听所有网卡
PORT = 50005      # 必须与 Pico 端一致

# 音频参数
SAMPLE_RATE = 16000
CHUNK_SIZE = 512  # 发送回传时的切片大小

# 临时文件名
TEMP_FILENAME = "last_echo.wav"

def process_and_save_wav(pcm_data, filename):
    """ 
    [核心逻辑更新]
    不再进行降噪滤波，仅使用位移操作将 32-bit 数据转换为 16-bit。
    """
    print(f"[Processing] Generating {filename} ...")
    
    # 1. 读取原始 32-bit 数据 (Little Endian)
    raw_data = np.frombuffer(pcm_data, dtype='<i4')
    
    if len(raw_data) == 0: 
        print("Error: Empty data received.")
        return False

    # 2. 核心处理：位移缩放
    # 你的测试代码证明右移 16 位是效果最好的
    shift_amount = 16 
    
    # 将 32-bit 映射到 16-bit 范围
    # 注意：直接位移，不经过滤波，保留原始底噪和细节
    scaled_data = raw_data >> shift_amount
    
    # 3. 安全削波 (Clipping) 防止爆音
    scaled_data = np.clip(scaled_data, -32768, 32767)
    
    # 4. 转为 int16 格式
    pcm_16bit = scaled_data.astype(np.int16)
    
    # 5. 保存文件
    wavfile.write(filename, SAMPLE_RATE, pcm_16bit)
    
    # 显示完整路径 (你的需求)
    full_path = os.path.abspath(filename)
    print(f"[Storage] Saved successfully to: {full_path}")
    return True

def stream_audio_back(target_ip, target_port, filename):
    """ 读取 WAV 文件并流式传回 Pico """
    if not os.path.exists(filename): 
        print(f"Error: File {filename} not found.")
        return

    print(f"[Playback] Streaming back to {target_ip}:{target_port} ...")
    
    wf = wave.open(filename, 'rb')
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    
    try:
        data = wf.readframes(CHUNK_SIZE // 2)
        
        while data:
            # 发送数据包
            sock.sendto(data, (target_ip, target_port))
            
            # 读取下一块
            data = wf.readframes(CHUNK_SIZE // 2)
            
            # === [流控的关键] ===
            # 防止电脑发送太快把 Pico 噎死
            # 16kHz * 16bit = 32000 bytes/sec
            # 512 bytes / 32000 ≈ 0.016s
            # 稍微留一点余量，0.015s 是经验值
            time.sleep(0.015) 

        # 发送结束标志
        sock.sendto(b'END_OF_STREAM', (target_ip, target_port))
        print("[Playback] Done.")
        
    except Exception as e:
        print(f"Stream Error: {e}")
    finally:
        sock.close()
        wf.close()

def start_server():
    server_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server_sock.bind((HOST, PORT))
    print(f"=== Echo Server Started on Port {PORT} ===")
    print("Logic: Record -> Bit Shift (No Filter) -> Auto Replay")
    print("------------------------------------------------")

    frames = []
    is_recording = False
    client_addr = None # 自动记录 Pico 的地址

    try:
        while True:
            # 接收数据
            data, addr = server_sock.recvfrom(4096)

            # 协议逻辑
            if b'START' in data:
                print(f"[{datetime.datetime.now().strftime('%H:%M:%S')}] Recording started from {addr}...")
                frames = []
                is_recording = True
                client_addr = addr # 记住是谁发来的，等会儿发回去
                
            elif b'STOP' in data:
                print("\nStop signal received.")
                is_recording = False
                
                if len(frames) > 0 and client_addr:
                    # 1. 合并数据
                    full_data = b''.join(frames)
                    
                    # 2. 处理并保存 (纯位移，无滤波)
                    # 使用 TEMP_FILENAME 以便循环覆盖，节省空间
                    if process_and_save_wav(full_data, TEMP_FILENAME):
                        
                        # 3. 立即回放 (Echo)
                        stream_audio_back(client_addr[0], client_addr[1], TEMP_FILENAME)
                        
                    print("--- Ready for next turn ---")
                else:
                    print("Warning: No audio data received.")

            elif is_recording:
                frames.append(data)
                # 每接收 50 个包打印一个点，证明数据在传输
                if len(frames) % 50 == 0:
                    print(".", end="", flush=True)

    except KeyboardInterrupt:
        print("\nServer stopped")
    finally:
        server_sock.close()

if __name__ == '__main__':
    start_server()